package taskqueue;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;
import redis.clients.jedis.Transaction;

import java.util.Set;
import java.util.Timer;
import java.util.TimerTask;
import java.util.UUID;


public class TaskQueue {

    private String idMessageMapKey;
    private String messageIdKey;
    private String noackMessageIdKey;

    private SnowFlake id = new SnowFlake(1L, 1L);
    private JedisPool pool;
    private MessageConsumer messageConsumer;

    private int messageTimeoutSecond = 10;

    public TaskQueue(String redisServer, String name, MessageConsumer messageConsumer) {
        this.pool = new JedisPool(new JedisPoolConfig(), redisServer);
        this.messageConsumer = messageConsumer;
        this.idMessageMapKey = "TaskQueue." + name + ".messagemap";
        this.messageIdKey = "TaskQueue." + name + ".queue";
        this.noackMessageIdKey = "TaskQueue." + name + ".noack";
    }

    public void add(String task) {
        try (Jedis jedis = pool.getResource()) {
            String id = UUID.randomUUID().toString();

            System.out.println("add message, id=" + id);

            Transaction transaction = jedis.multi();
            transaction.hset(idMessageMapKey, id, task);
            transaction.zadd(messageIdKey, messageOrder(), id);
            transaction.exec();
        }
    }

    private double messageOrder() {
        return (double) id.nextId();
    }

    public void start() {
        new Timer().schedule(new TimerTask() {
            @Override
            public void run() {
                TaskQueue.this.noaskMessage();
            }
        }, 1000, 1000);
        new Consumer(this).start();
    }

    public void noaskMessage() {
        try (Jedis jedis = pool.getResource()) {

            while (true) {
                Set<String> ids = jedis.zrangeByScore(noackMessageIdKey, 0.0D, Double.MAX_VALUE, 0, 1);
                if (!ids.isEmpty()) {
                    String id = ids.iterator().next();
                    String value = jedis.set(id, "locked", "NX", "EX", 10);
                    if ("OK".equals(value)) {
                        System.out.println("no ack message, id=" + id);

                        Transaction transaction = jedis.multi();
                        transaction.zrem(noackMessageIdKey, id);
                        transaction.zadd(messageIdKey, messageOrder(), id);
                        transaction.expire(id, 0);
                        transaction.exec();
                    }
                }
            }
        }
    }

    private void startHandleMessage() {
        while (true) {
            Message task = this.get();
            if (task != null) {
                try {
                    this.messageConsumer.handle(task.getBody());
                    this.afterHandleMessage(task);
                } catch (Exception e) {
                    e.printStackTrace();
                    // 不支持重试,直接丢弃该消息
                    this.afterHandleMessage(task);
                }
            } else {
                try {
                    Thread.sleep(100);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    public Message get() {
        try (Jedis jedis = pool.getResource()) {
            Set<String> ids = jedis.zrangeByScore(messageIdKey, 0.0D, Double.MAX_VALUE, 0, 1);
            if (!ids.isEmpty()) {
                String id = ids.iterator().next();
                // lock
                String value = jedis.set(id, "locked", "NX", "EX", messageTimeoutSecond);
                if ("OK".equals(value)) {
                    String messageBody = jedis.hget(idMessageMapKey, id);

                    Transaction transaction = jedis.multi();
                    transaction.zrem(messageIdKey, id);
                    transaction.zadd(noackMessageIdKey, messageOrder(), id);
                    transaction.exec();

                    if(messageBody != null) {
                        // messageBody等于null,该message已经被消费
                        Message message = new Message();
                        message.setId(id);
                        message.setBody(messageBody);
                        return message;
                    }
                }
            }

            return null;
        }
    }

    public void afterHandleMessage(Message task) {
        try (Jedis jedis = pool.getResource()) {
            System.out.println("ack message id=" + task.getId());
            Transaction transaction = jedis.multi();
            transaction.zrem(noackMessageIdKey, task.getId());
            transaction.hdel(idMessageMapKey, task.getId());
            transaction.exec();
        }
    }

    private static class Message {

        private String id;
        private String body;

        public String getId() {
            return id;
        }

        public void setId(String id) {
            this.id = id;
        }

        public String getBody() {
            return body;
        }

        public void setBody(String body) {
            this.body = body;
        }
    }

    private static class Consumer extends Thread {

        TaskQueue queue;

        public Consumer(TaskQueue queue) {
            this.queue = queue;
        }

        @Override
        public void run() {
            queue.startHandleMessage();
        }
    }


}
